"""
Utility to draw a Gantt chart for pointer-net results.

Provides a single function `draw_gantt_from_results(results, g, out_path, ...)` which accepts
- results: list-like of routes (several common formats supported)
- g: graph-like object with `ndata['available_window']` and optional `ndata['goods']` and `num_nodes()`
- out_path: full path where the PNG will be written

The function is defensive about input formats: routes can be
- list of ints -> each step is a single time unit starting at 0
- list of (node, start, end) tuples -> used as-is
- nested or dict formats (best-effort flattening)

It will also draw per-node available windows if `g.ndata['available_window']` is available
(as an int bitmask per node). If there are overlapping time intervals for a single vehicle,
those intervals are adjusted to be sequential (so one vehicle isn't at multiple nodes at the same
exact time unit).

This file used to contain a standalone script; it's been refactored to expose the function
for programmatic use from `run_training.py`.
"""

from typing import Any, List, Sequence, Tuple, Optional
import os
import math
import logging

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

# Optional heavy dependency imports are done lazily inside functions to keep import fast.


def _ensure_list(x):
    if x is None:
        return []
    if isinstance(x, (list, tuple)):
        return list(x)
    return [x]


def _parse_route_entry(entry) -> List[int]:
    """Try to flatten a route entry into a list of node ints.
    If entry already looks like [(node, start, end), ...] this is handled by caller.
    """
    out = []
    if entry is None:
        return out
    if isinstance(entry, dict):
        # common key names
        for k in ('route', 'path', 'tour'):
            if k in entry:
                return _parse_route_entry(entry[k])
        # fallback: flatten values
        for v in entry.values():
            out.extend(_parse_route_entry(v))
        return out
    if isinstance(entry, (list, tuple)):
        # if each element is int-like -> return them
        if all(isinstance(x, int) for x in entry):
            return list(entry)
        # otherwise flatten recursively
        for item in entry:
            if isinstance(item, (list, tuple)) and len(item) == 3 and all(isinstance(i, (int, float)) for i in item):
                # detected (node, start, end) pattern -> stop and signal caller
                return list(entry)
            try:
                out.extend(_parse_route_entry(item))
            except Exception:
                pass
        return out
    # single int-like
    try:
        ival = int(entry)
        return [ival]
    except Exception:
        return []


def _available_window_to_intervals(val: int, bitwidth: Optional[int] = None, bit_order: str = 'msb') -> List[Tuple[int, int]]:
    """Convert an integer bitmask into a list of (start, end) intervals in time units.
    By default this interprets the printed binary strings (format(...)) where the leftmost
    character is the earliest time unit (MSB-first). Use `bit_order='lsb'` to interpret
    bit 0 (least-significant bit) as time 0.
    """
    if val is None:
        return []
    if bitwidth is None:
        bitwidth = max(1, val.bit_length())
    intervals = []
    cur_start = None
    for t in range(bitwidth):
        # map t (time index starting at 0) to the corresponding bit position
        if bit_order == 'msb':
            bit_pos = bitwidth - 1 - t
        else:
            bit_pos = t
        if ((val >> bit_pos) & 1) == 1:
            if cur_start is None:
                cur_start = t
        else:
            if cur_start is not None:
                intervals.append((cur_start, t))
                cur_start = None
    if cur_start is not None:
        intervals.append((cur_start, bitwidth))
    return intervals


def draw_gantt_from_results(
    results: Sequence[Any],
    g: Any,
    out_path: str,
    num_vehicles: Optional[int] = None,
    num_nodes: Optional[int] = None,
    max_time: Optional[int] = None,
    time_unit: float = 1.0,
    colors: Optional[List[str]] = None,
    figsize: Tuple[int, int] = (12, 8),
    show_available_windows: bool = True,
    bit_order: str = 'msb',
) -> str:
    """Draw a Gantt chart from pointer-network `results` and a graph-like `g`.

    Parameters
    - results: sequence where each element is one vehicle's route. Supported per-vehicle formats:
        * [node0, node1, ...]  (integers): each node occupies one `time_unit` sequentially starting at 0
        * [(node, start, end), ...] : explicit intervals (start,end) in same units used for x-axis
        * dict/other nested structures: best-effort flattening
    - g: graph-like object with:
        * g.num_nodes() -> int OR pass num_nodes explicitly
        * g.ndata['available_window'] -> 1-D tensor/list of ints (bitmask) optional
    - out_path: where to write the PNG file (parent directories will be created)
    - num_vehicles: if None, inferred from results length
    - num_nodes: if None, inferred from g.num_nodes() if possible, otherwise from max node index in results
    - max_time: if None, inferred from available_window bitwidth or from routes
    - time_unit: multiplier for discrete steps -> real x-axis units
    - colors: list of colors for vehicles
    - figsize: matplotlib figure size
    - show_available_windows: draw light background bars that indicate the available window per node

    Returns: the out_path written (for convenience)
    """
    # defensive imports
    try:
        import torch
    except Exception:
        torch = None

    # normalize results to list of per-vehicle route representations
    rlist = list(results) if results is not None else []
    if num_vehicles is None:
        num_vehicles = len(rlist)
    if colors is None:
        base_colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple']
        colors = [base_colors[i % len(base_colors)] for i in range(num_vehicles)]

    # infer num_nodes
    inferred_num_nodes = None
    if num_nodes is None:
        # try g.num_nodes()
        try:
            inferred_num_nodes = int(g.num_nodes())
        except Exception:
            inferred_num_nodes = None
    else:
        inferred_num_nodes = num_nodes

    # Build intervals per vehicle: List[List[(node, start, end)]]
    vehicle_intervals: List[List[Tuple[int, float, float]]] = []

    overall_max_time = 0
    max_node_index = -1

    for v_idx in range(num_vehicles):
        if v_idx < len(rlist):
            entry = rlist[v_idx]
        else:
            entry = []
        # detect explicit interval format
        use_explicit = False
        if isinstance(entry, (list, tuple)) and len(entry) > 0 and all(isinstance(x, (list, tuple)) and len(x) == 3 for x in entry):
            use_explicit = True
        intervals = []
        if use_explicit:
            for node, start, end in entry:
                try:
                    node_i = int(node)
                    s = float(start)
                    e = float(end)
                except Exception:
                    continue
                intervals.append((node_i, s, e))
                overall_max_time = max(overall_max_time, e)
                max_node_index = max(max_node_index, node_i)
        else:
            # try flattening to ints
            flattened = _parse_route_entry(entry)
            # if flatten returned a list that itself contains tuples (detected earlier), it would have been handled
            # assign sequential times: step k occupies [k*time_unit, (k+1)*time_unit)
            for step_idx, node in enumerate(flattened):
                try:
                    node_i = int(node)
                except Exception:
                    continue
                s = step_idx * time_unit
                e = s + time_unit
                intervals.append((node_i, s, e))
                overall_max_time = max(overall_max_time, e)
                max_node_index = max(max_node_index, node_i)
        vehicle_intervals.append(intervals)

    # infer num_nodes from max index if still missing
    if inferred_num_nodes is None:
        if max_node_index >= 0:
            inferred_num_nodes = max_node_index + 1
        else:
            inferred_num_nodes = 0

    # determine max_time
    if max_time is None:
        # prefer available_window bitwidth if present
        bitwidth = None
        aw_values = None
        try:
            aw_tensor = None
            if hasattr(g, 'ndata') and isinstance(g.ndata, dict):
                aw_tensor = g.ndata.get('available_window', None)
            else:
                try:
                    aw_tensor = g.ndata['available_window']
                except Exception:
                    aw_tensor = None
            if aw_tensor is not None:
                if torch is not None and isinstance(aw_tensor, torch.Tensor):
                    aw_cpu = aw_tensor.detach().cpu().view(-1)
                    aw_values = [int(x.item()) for x in aw_cpu]
                else:
                    try:
                        aw_values = [int(x) for x in list(aw_tensor)]
                    except Exception:
                        # single int
                        try:
                            aw_values = [int(aw_tensor)]
                        except Exception:
                            aw_values = None
                if aw_values:
                    max_val = max(aw_values)
                    bitwidth = max(1, max_val.bit_length())
        except Exception:
            aw_values = None
            bitwidth = None
        if bitwidth is not None:
            max_time = bitwidth
        else:
            max_time = max(overall_max_time, 1)
    # ensure sensible integer for plotting ticks
    try:
        max_time = int(math.ceil(max_time))
    except Exception:
        max_time = 10

    # If show_available_windows, compute per-node intervals
    node_windows = None
    if show_available_windows:
        node_windows = []
        try:
            aw_tensor = None
            if hasattr(g, 'ndata') and isinstance(g.ndata, dict):
                aw_tensor = g.ndata.get('available_window', None)
            else:
                try:
                    aw_tensor = g.ndata['available_window']
                except Exception:
                    aw_tensor = None
            if aw_tensor is not None:
                if torch is not None and isinstance(aw_tensor, torch.Tensor):
                    aw_cpu = aw_tensor.detach().cpu().view(-1)
                    aw_values = [int(x.item()) for x in aw_cpu]
                else:
                    aw_values = [int(x) for x in list(aw_tensor)]
                # determine bitwidth
                bitwidth = max(1, max((v.bit_length() for v in aw_values), default=0))
                bitwidth = max(bitwidth, max_time)
                for v in aw_values:
                    node_windows.append(_available_window_to_intervals(v, bitwidth, bit_order=bit_order))
            else:
                node_windows = []
        except Exception:
            node_windows = []

    # Fix overlapping intervals per vehicle: ensure sequential
    adjusted_vehicle_intervals: List[List[Tuple[int, float, float]]] = []
    for v_idx, intervals in enumerate(vehicle_intervals):
        # sort by start
        iv = sorted(intervals, key=lambda x: x[1])
        adj = []
        cur_end = None
        for node, s, e in iv:
            if cur_end is None:
                adj.append((node, s, e))
                cur_end = e
            else:
                if s < cur_end:
                    # shift forward so it starts at cur_end
                    duration = max(0.0001, e - s)
                    ns = cur_end
                    ne = ns + duration
                    logging.warning("Adjusting overlapping interval for vehicle %s: node %s from (%s,%s) -> (%s,%s)", v_idx, node, s, e, ns, ne)
                    adj.append((node, ns, ne))
                    cur_end = ne
                else:
                    adj.append((node, s, e))
                    cur_end = e
        adjusted_vehicle_intervals.append(adj)
        for (_, _, e) in adj:
            overall_max_time = max(overall_max_time, e)

    # create figure
    fig, ax = plt.subplots(figsize=figsize)

    # y positions: use node index as integer rows
    y_ticks = list(range(inferred_num_nodes))

    # draw available windows as light rectangles per node
    if node_windows:
        for node_idx, intervals in enumerate(node_windows):
            for (s, e) in intervals:
                ax.broken_barh([(s, e - s)], (node_idx - 0.45, 0.9), facecolors=('0.95'), edgecolors=None)

    # draw vehicle bars
    for v_idx, intervals in enumerate(adjusted_vehicle_intervals):
        color = colors[v_idx % len(colors)]
        for (node, s, e) in intervals:
            if node < 0 or node >= inferred_num_nodes:
                # skip out-of-range nodes but log
                logging.debug("Skipping out-of-range node %s for vehicle %s", node, v_idx)
                continue
            # shift small vertical offset per vehicle to avoid full overlap
            offset = (v_idx - (num_vehicles - 1) / 2.0) * 0.2
            ax.broken_barh([(s, e - s)], (node - 0.3 + offset, 0.2), facecolors=color, edgecolor='k', linewidth=0.3, alpha=0.9)

    ax.set_xlabel('Time')
    ax.set_ylabel('Node')
    ax.set_yticks(y_ticks)
    ax.set_yticklabels([str(i) for i in y_ticks])
    ax.set_title('Gantt Chart of Vehicle Routing')
    ax.set_xlim(0, max_time)
    ax.set_ylim(-1, inferred_num_nodes)

    # legend
    handles = [mpatches.Patch(color=colors[i % len(colors)], label=f'Vehicle {i+1}') for i in range(num_vehicles)]
    ax.legend(handles=handles, loc='upper right')

    plt.tight_layout()

    # write file
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    fig.savefig(out_path)
    plt.close(fig)
    logging.info('Gantt chart saved to %s', out_path)
    return out_path


# Keep a small CLI example for debugging; does not run on import
if __name__ == '__main__':
    # lightweight demo if executed directly
    try:
        import torch
        class DummyG:
            def __init__(self, n):
                self.ndata = {}
                self._n = n
            def num_nodes(self):
                return self._n
        g = DummyG(20)
        # sample available window: allow times 0-4 and 6-9 for node 0
        # This demo uses MSB-left binary representation: '1011101100' etc would be interpreted left->right as time 0..9
        g.ndata['available_window'] = [0b0000011111] * 20
        results = [ [0,5,8], [1,3,7], [2,4,6] ]
        out = os.path.join(os.path.dirname(__file__), 'gantt_demo.png')
        draw_gantt_from_results(results, g, out, bit_order='msb')
    except Exception as e:
        logging.error('Demo failed: %s', e)
